#######################################################

library(ggplot2)
library(Seurat)
library(magrittr)
library(pheatmap)
library(parallel)
library(dplyr)
library(pheatmap)
library(optparse)

##########################################################################################
option_list <- list(
    make_option(c("--rna_file"), type = "character"),
    make_option(c("--type"), type = "character"),
    make_option(c("--divide"), type = "character"),
    make_option(c("--scriptPath"), type = "character"),
    make_option(c("--pct"), type = "character"),
    make_option(c("--logfc"), type = "character"),
    make_option(c("--cpu"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 单细胞rna文件
    rna_file <- "~/20231121_singleMuti/results/qc_atac_v3/somatic/testis_combined.annotationCellType.qc.Rdata"

    ## 既往研究整理的代码
    scriptPath <- "~/20231121_singleMuti/scripts/qc_atac_v3/germ/scScalpChromatin"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/celltype_plot/diff_expression/germ"

    ## 
    type <- "germ"

    ## 生殖细胞分成几类
    divide <- 3

    ## pct
    pct <- 0.25

    ## fc 
    logfc <- 1

    ## cpu
    cpu <- 10
}

###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

rna_file <- opt$rna_file
scriptPath <- opt$scriptPath
divide <- as.numeric(opt$divide)
pct <- as.numeric(opt$pct)
logfc <- as.numeric(opt$logfc)
cpu <- as.numeric(opt$cpu)
out_path <- opt$out_path
type <- opt$type

dir.create( out_path , recursive = T )

###########################################################################################
## 读数据
load(rna_file)
## 使用不填补的
DefaultAssay(scrnat) <- "RNA"

##########################################################################################
## 已发表文献写好的脚本
source(paste0(scriptPath, "/plotting_config.R"))
source(paste0(scriptPath, "/misc_helpers.R"))
source(paste0(scriptPath, "/matrix_helpers.R"))
source(paste0(scriptPath, "/archr_helpers.R"))
source(paste0(scriptPath, "/GO_wrappers.R"))

###########################################################################################
## 细胞顺序
grp_order2 = c("SSC",
"Differenting&Differented SPG",
"Leptotene",
"Zygotene",
"Patchytene",
"Diplotene",
"Early stage of spermatids",
"Round&ElongateS.tids",
"Sperm",
"Leydig cells",
"Myoid cells",
"Pericytes",
"Sertoli cells",
"Endothelial cells",
"NKT cells",
"Macrophages")

SSC_SPG <- c("SSC" , "Differenting&Differented SPG")
SPC <- c("Leptotene" , "Zygotene", "Patchytene",
    "Diplotene" , "Early stage of spermatids")
SPT <- c("Round&ElongateS.tids" , "Sperm")

###########################################################################################

if(divide == 3){
  ## 合并细胞为三大类
  scrnat$cell_type <- ifelse(scrnat$cell_type %in% SSC_SPG , "SSC_SPG" , scrnat$cell_type)
  scrnat$cell_type <- ifelse(scrnat$cell_type %in% SPC , "SPC" , scrnat$cell_type)
  scrnat$cell_type <- ifelse(scrnat$cell_type %in% SPT , "SPT" , scrnat$cell_type)
  scrnat$cell_type <- factor( scrnat$cell_type , levels = c("SSC_SPG" , "SPC" , "SPT") , order = T )
}else{
  scrnat$cell_type <- factor( scrnat$cell_type , levels = grp_order2[grp_order2 %in% unique(scrnat$cell_type)] , order = T )
}

## 确认排序
Idents(scrnat) <- scrnat$cell_type

germ = scrnat
cellNames <- levels(germ)

###########################################################################################
## 富集分析的函数
pathway_enrich <- function(all_genes = all_genes , clust_genes = clust_genes){
  upGO <- rbind(
    calcTopGo(all_genes, interestingGenes=clust_genes, nodeSize=5, ontology="BP") 
    )

  upGO <- upGO[order(as.numeric(upGO$pvalue), decreasing=FALSE),]

  ## 构造GO对象
  geneList <- factor(as.integer(all_genes %in% clust_genes))
  names(geneList) <- all_genes
  GOdata <- suppressMessages(new(
    "topGOdata",
    ontology = "BP",
    allGenes = geneList,
    annot = annFUN.org, mapping = "org.Hs.eg.db", ID = "symbol",
    nodeSize = 5
  ))
  ## 提取每个通路中感兴趣的基因
  gene_in_pathway <- sapply(upGO$GO.ID, function(x)
    {
      genes <- genesInTerm(GOdata, x)
      # myGenes is the queried gene list
      paste0(genes[[1]][genes[[1]] %in% clust_genes] , collapse = "," )
    })
  
  upGO$gene_in_pathway <- gene_in_pathway

  return(upGO)
}

###########################################################################################
############################################
## 邻近的细胞间两两做差异表达
dat_diff <- bind_rows(mclapply(2:length(cellNames),function(k){
  print(k)
  celltype1 <- cellNames[[k]]
  celltype2 <- cellNames[[k-1]]
  
  ## 不卡foldchange
  ## 后期的细胞对前期的foldchange
  tmp <- FindMarkers(germ, assay="RNA", 
      ident.1 = cellNames[[k]], ident.2 = cellNames[[k-1]], only.pos=F, 
      min.pct = pct , logfc.threshold = 0 )
  tmp$cell_compare <- paste0( celltype1 , "-" , celltype2 )

  return(tmp)

},mc.cores=cpu))

dat_diff$gene <- sapply( strsplit(rownames(dat_diff) , "[.][.][.]") , "[" , 1 )
out_file <- paste0( out_path , "/" , "two_adjacent_celltype." , type , "." , "pct_" , pct , ".tsv" )
write.table( dat_diff , out_file , row.names = F , quote = F , sep ="\t" )

############################################
## 通路富集分析比较慢
GOresults <- bind_rows(mclapply(unique(dat_diff$cell_compare),function(cellT){
  message(sprintf("Running GO enrichments on %s...", cellT))
  all_genes <- rownames(scrnat)

  ## 提取显著上调的基因
  clust_genes <- subset( dat_diff , cell_compare==cellT & avg_log2FC > logfc & p_val_adj < 0.05 )$gene
  up_pathway <- pathway_enrich(all_genes = all_genes , clust_genes = clust_genes)
  up_pathway$up_down <- "up"

  ## 提取显著下调的基因
  clust_genes <- subset( dat_diff , cell_compare==cellT & avg_log2FC < -logfc & p_val_adj < 0.05 )$gene
  down_pathway <- pathway_enrich(all_genes = all_genes , clust_genes = clust_genes)  
  down_pathway$up_down <- "down"

  tmp <- rbind( up_pathway , down_pathway )
  tmp$cell_compare <- cellT

  return(tmp)

},mc.cores=cpu))

out_file <- paste0( out_path , "/" , "two_adjacent_celltype." , type , "." , "pct_" , pct, ".logfc_" , logfc , ".GOBP.tsv" )
write.table( GOresults , out_file , row.names = F , quote = F , sep ="\t" )

###########################################################################################
############################################
## 一类细胞和其它所有细胞做差异表达
## 只看高表达
dat_diff <- bind_rows(mclapply(1:length(cellNames),function(k){
  print(k)
  celltype1 <- cellNames[[k]]
  
  ## 不卡foldchange
  ## 后期的细胞对前期的foldchange
  tmp <- FindMarkers(germ, assay="RNA", 
      ident.1 = cellNames[[k]], only.pos=F, 
      min.pct = pct , logfc.threshold = 0 )
  tmp$cell_compare <- celltype1

  return(tmp)

},mc.cores=cpu))

dat_diff$gene <- sapply( strsplit(rownames(dat_diff) , "[.][.][.]") , "[" , 1 )
out_file <- paste0( out_path , "/" , "one_vs_other." , type , "." , "pct_" , pct , ".tsv" )
write.table( dat_diff , out_file , row.names = F , quote = F , sep ="\t" )

############################################
## 通路富集
GOresults <- bind_rows(mclapply(unique(dat_diff$cell_compare),function(cellT){
  message(sprintf("Running GO enrichments on %s...", cellT))
  all_genes <- rownames(scrnat)

  ## 提取显著上调的基因
  clust_genes <- subset( dat_diff , cell_compare==cellT & avg_log2FC > logfc & p_val_adj < 0.05 )$gene
  up_pathway <- pathway_enrich(all_genes = all_genes , clust_genes = clust_genes)
  up_pathway$up_down <- "up"

  tmp <- up_pathway
  tmp$cell_compare <- cellT

  return(tmp)

},mc.cores=cpu))

out_file <- paste0( out_path , "/" , "one_vs_other." , type , "." , "pct_" , pct, ".logfc_" , logfc , ".GOBP.tsv" )
write.table( GOresults , out_file , row.names = F , quote = F , sep ="\t" )

############################################
## 差异基因热图
## 构造表达矩阵
if(1!=1){
  use_gene <- unique(subset( dat_diff , avg_log2FC > logfc & p_val_adj < 0.05 )$gene)
  sco_exp <- sapply(grp_order,function(x){
      sapply(unique(use_gene),function(y){
          mean(as.numeric(as.vector(scrnat@assays$RNA@data[y,which(scrnat$cell_type==x)])))
      })
  })
  sco_exp <- sco_exp[which(rowSums(sco_exp)>0),]
  sco_exp <- sco_exp[,grp_order]

  ## 画图
  p <- pheatmap(sco_exp,scale = "row",
      show_rownames = F ,show_colnames = T, 
      #cutree_cols=7,cutree_rows = 7,
      cluster_rows = T , cluster_cols = F , 
      clustering_method = "ward.D2",
      color = colorRampPalette(c("blue", "white", "firebrick3"))(100),
      cellwidth = 15, cellheight = 3)

  out_file <- paste0( out_path , "/" , "one_vs_other." , type , "." , "pct_" , pct , ".logfc_" , logfc , ".pdf" )
  pdf(out_file , height = 15 , width = 8)
  p
  dev.off()
}